import chromadb.utils.embedding_functions as embedding_functions
import chromadb
import json

ollama_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="sk-lOnR4erWFrXwpBqW24E293122a0347AdA983D7E99e0512B1",
                api_base="http://60.204.226.75:32100/v1",
                model_name="mxbai-embed-large"
            )
documents=[]
# documents=["灯杆/单灯/方案信息相关表：{'灯源状态': 'light.l_coplog','巡测历史记录': 'light.l_his_coplog'}",
#            "视频设备信息表：{'视频服务器管理': 'device.device_camera_server','视频设备管理': 'device.device_camera'}",
#            "公共广播信息表：{'公共广播信息表': 'device.device_broadcast'}"]

chroma_client = chromadb.PersistentClient(path="./db")
collection = chroma_client.get_or_create_collection(name="my_demo04_collection", embedding_function=ollama_ef)

for i, d in enumerate(documents):
  embedding = ollama_ef(d)
  collection.upsert(
    ids=[str(i)],
    embeddings=embedding,
    documents=[d]
  )

query_results = collection.query(query_embeddings=ollama_ef("查询广播点的灯杆信息"))
print(f"查询结果：{query_results}")
